# Imports for gradio
import gradio as gr
import random
from http import HTTPStatus

# Imports for dashscope
import dashscope
import json

# Imports for model
from resnet import ResNet1d
import tqdm
import h5py
import torch
import os
import json
import numpy as np
import argparse
from warnings import warn
import pandas as pd

import sys

# 调用阿里千问API
def call_ai_with_messages(content):
    messages = [
        {'role': 'system', 'content': '你是一名人工智能辅助医生，具备丰富的心血管疾病知识和经验。您的任务是使用最新的人工智能工具，为心电图的诊断提供帮助，结合患者的年龄、性别、心电图预测年龄，来给出合理的护理建议。'},
        {'role': 'user', 'content': '%s\n请根据这些信息，给出适当的护理建议，护理建议中，包括三项主要建议，三项辅助建议。' % content}
    ]
    response = dashscope.Generation.call(
        dashscope.Generation.Models.qwen_turbo,
        messages=messages,
        # set the random seed, optional, default to 1234 if not set
        seed=random.randint(1, 10000),
        result_format='result_format',  # set the result to be "message" format.
        max_tokens = 500
    )
    if response.status_code == HTTPStatus.OK:
        print(type(response))
        print("回复：", response['output']['text'])
        return response['output']['text']
    else:
        print('Request id: %s, Status code: %s, error code: %s, error message: %s' % (
            response.request_id, response.status_code,
            response.code, response.message
        ))
        return "暂无护理建议"

# graid API接口
def predict_process(name, sex, age, tip, file, memo, history):
    history = history or []

    if len(name) < 1:
        message = None
        response = "请输入患者编号"
    elif sex is None:
        message = None
        response = "请输入患者性别"
    elif age < 1:
        message = None
        response = "请输入患者年龄"
    elif file is None:
        message = None
        response = "请提交患者心电图检测数据"
    else:
        sex_name = "男性" if sex == '男' else "女性"
        if len(memo) >0:
            message = f"患者信息：编号：{name}，性别：{sex}，年龄：{age}岁，补充信息：{memo}"
        else:
            message = f"患者信息：编号：{name}，性别：{sex}，年龄：{age}岁"

        # 获取心电图年龄预测结果
        print("心电图数据文件：%s" % file.name)
        predict_ages = predict_age_from_traces(file.name)
        if not predict_ages or len(predict_ages)<1:
            response = '预测结果：数据有误，请重新上传心电图数据'
        else:
            predict_age = round(predict_ages[0], 1)

            if predict_age - age > 8:
                predict_die = '较高'
            elif predict_age - age > -8:
                predict_die = '较低'
            else:
                predict_die = '无法判断'

            if tip:
                # 获取护理建议
                predict_age_diff = predict_age - age
                if len(memo)>0:
                    request_content = f"患者性别：{sex_name}，实际年龄：{age}岁，心电图年龄：{predict_age}岁，心电图年龄比实际年龄大${predict_age_diff}岁，死亡风险：{predict_die}，患者补充信息：{memo}"
                else:
                    request_content = f"患者性别：{sex_name}，实际年龄：{age}岁，心电图年龄：{predict_age}岁，心电图年龄比实际年龄大${predict_age_diff}岁，死亡风险：{predict_die}"

                help_message = call_ai_with_messages(request_content)
                response = f"预测结果：\n心电图年龄：{predict_age}岁，死亡风险：{predict_die}\n\n\n护理建议：\n{help_message}"
            else:
                response = f"预测结果：\n心电图年龄：{predict_age}岁，死亡风险：{predict_die}"

    history.append((message, response))
    return history

# 心电图年龄预测
def predict_age_from_traces(path_to_traces, traces_dset='tracings', ids_dset=None, batch_size=8):
    # Get traces
    print("Get traces from %s" % path_to_traces)
    ff = h5py.File(path_to_traces, 'r')
    traces = ff[traces_dset]
    n_total = len(traces)
    if ids_dset:
        ids = ff[ids_dset]
    else:
        ids = range(n_total)
    # Get dimension
    predicted_age = np.zeros((n_total,))

    n_total, n_samples, n_leads = traces.shape
    n_batches = int(np.ceil(n_total/batch_size))
    # Compute gradients
    predicted_age = np.zeros((n_total,))
    end = 0
    for i in tqdm.tqdm(range(n_batches)):
        start = end
        end = min((i + 1) * batch_size, n_total)
        with torch.no_grad():
            x = torch.tensor(traces[start:end, :, :]).transpose(-1, -2)
            x = x.to(device, dtype=torch.float32)
            y_pred = model(x)
        predicted_age[start:end] = y_pred.detach().cpu().numpy().flatten()

    return predicted_age

if __name__ == "__main__":
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('mdl', type=str,
                        help='folder containing model.')
    parser.add_argument('--batch_size', type=int, default=8,
                        help='number of exams per batch.')
    parser.add_argument('--traces_dset', default='tracings',
                         help='traces dataset in the hdf5 file.')
    parser.add_argument('--ids_dset',
                         help='ids dataset in the hdf5 file.')
    args, unk = parser.parse_known_args()
    # Check for unknown options
    if unk:
        warn("Unknown arguments:" + str(unk) + ".")

    # 计算设备检查
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # 模型加载
    print("加载模型")
    ckpt = torch.load(os.path.join(args.mdl, 'model.pth'), map_location=lambda storage, loc: storage)
    # Get config
    config = os.path.join(args.mdl, 'config.json')
    with open(config, 'r') as f:
        config_dict = json.load(f)
    # Get model
    N_LEADS = 12
    model = ResNet1d(input_dim=(N_LEADS, config_dict['seq_length']),
                     blocks_dim=list(zip(config_dict['net_filter_size'], config_dict['net_seq_lengh'])),
                     n_classes=1,
                     kernel_size=config_dict['kernel_size'],
                     dropout_rate=config_dict['dropout_rate'])
    # load model checkpoint
    model.load_state_dict(ckpt["model"])
    model = model.to(device)

    # Evaluate on test data
    model.eval()

    # gradio服务器启动
    print("启动WEB服务")
    css = """
    #stitle {text-align:center}
    #schatbot {height:500px;}
    #sfile {height:150px;}
    """

    with gr.Blocks(css=css) as webui:
        gr.Markdown(
            """
            # 心电图年龄及死亡风险预测系统

            ### 根据提交的12导联心电图数据，预测其心电图年龄，并与患者年龄结合预测死亡风险
            """,elem_id="stitle")

        with gr.Row():
            # 定义chat机器人
            chatbot = gr.Chatbot(elem_id="schatbot", label="预测结果", height="")

            # 定义患者信息输入
            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        in_name = gr.Textbox(lines=1,placeholder='请输入编号', label="患者编号")
                        in_sex = gr.Radio(["男", "女"], label="患者性别")
                        in_tip = gr.Checkbox(label="护理建议", info="是否通过AI提供护理建议")
                    in_age = gr.Slider(0, 120, label="患者年龄")
                    in_file = gr.File(label="心电图检测数据", elem_id="sfile")
                    in_memo = gr.Textbox(lines=3,placeholder='请输入患者补充信息', label="补充信息")

        with gr.Row():
        # 定义按钮
            clear = gr.ClearButton([chatbot])
            btn = gr.Button("提交")
            btn.click(
                fn=predict_process,
                inputs=[in_name,in_sex,in_age,in_tip,in_file,in_memo,chatbot],
                outputs=[chatbot]
            )

    webui.title = "心电图年龄及死亡风险预测系统"
    webui.launch(server_name="0.0.0.0", server_port=7860, share=False)